"""
Type conversion utilities with uppercase function names.

This module provides robust type conversion functions following mainstream
deep learning framework patterns (PyTorch Lightning, JAX, HuggingFace).
"""

import weakref
import threading
from typing import Any, Union, Optional, Tuple, List, Sequence
from collections.abc import Iterable

import torch
import torch.nn as nn
import numpy as np

from continuallearning.utils.logging import get_logger

logger = get_logger(__name__)

# Type aliases following JAX/PyTorch Lightning patterns
TENSOR_LIKE = Union[torch.Tensor, np.ndarray, List, Tuple, float, int]
TENSOR_INPUT = Union[torch.Tensor, Tuple[torch.Tensor, ...], List[torch.Tensor]]
MODEL_OUTPUT = Union[torch.Tensor, Tuple[torch.Tensor, ...]]
NUMERIC_SCALAR = Union[int, float, np.number]
ARRAY_LIKE = Union[np.ndarray, List, Tuple, torch.Tensor]


def TO_TENSOR(
    data: TENSOR_LIKE,
    dtype: Optional[torch.dtype] = None,
    device: Optional[torch.device] = None,
    requires_grad: bool = False,
) -> torch.Tensor:
    """
    Convert various input types to PyTorch tensor.

    This function handles the most common tensor conversion scenarios
    with comprehensive error handling and type validation.

    Args:
        data: Input data to convert to tensor
        dtype: Target tensor dtype
        device: Target device for tensor
        requires_grad: Whether tensor requires gradient

    Returns:
        PyTorch tensor

    Raises:
        TypeError: If input type is not supported
        ValueError: If conversion fails due to invalid data

    Examples:
        >>> TO_TENSOR([1, 2, 3])
        >>> TO_TENSOR(np.array([1.0, 2.0]))
        >>> TO_TENSOR(3.14, dtype=torch.float32)
    """
    if data is None:
        raise ValueError("Cannot convert None to tensor")

    try:
        # Already a tensor - handle device/dtype conversion
        if torch.is_tensor(data):
            result = data
            if dtype is not None and result.dtype != dtype:
                result = result.to(dtype=dtype)
            if device is not None and result.device != device:
                result = result.to(device=device)
            if requires_grad != result.requires_grad:
                result = result.requires_grad_(requires_grad)
            return result

        # NumPy array
        if isinstance(data, np.ndarray):
            if data.dtype == np.object_:
                raise ValueError("Cannot convert object array to tensor")
            result = torch.from_numpy(data)

        # Scalar numeric types
        elif isinstance(data, (int, float, np.number)):
            result = torch.tensor(data)

        # Sequence types (list, tuple)
        elif isinstance(data, (list, tuple)):
            if len(data) == 0:
                raise ValueError("Cannot convert empty sequence to tensor")
            result = torch.tensor(data)

        # Other iterables
        elif isinstance(data, Iterable):
            try:
                result = torch.tensor(list(data))
            except (ValueError, TypeError) as e:
                raise TypeError(
                    f"Cannot convert iterable of type {type(data)} to tensor: {e}"
                )

        else:
            raise TypeError(f"Unsupported type for tensor conversion: {type(data)}")

        # Apply optional transformations
        if dtype is not None:
            result = result.to(dtype=dtype)
        if device is not None:
            result = result.to(device=device)
        if requires_grad:
            result = result.requires_grad_(requires_grad)

        return result

    except Exception as e:
        logger.error(f"Failed to convert {type(data)} to tensor: {e}")
        raise


def TO_NUMPY(data: TENSOR_LIKE) -> np.ndarray:
    """
    Convert various input types to NumPy array.

    Args:
        data: Input data to convert to numpy array

    Returns:
        NumPy array

    Raises:
        TypeError: If input type is not supported

    Examples:
        >>> TO_NUMPY(torch.tensor([1, 2, 3]))
        >>> TO_NUMPY([1.0, 2.0, 3.0])
    """
    if data is None:
        raise ValueError("Cannot convert None to numpy array")

    try:
        # PyTorch tensor
        if torch.is_tensor(data):
            if data.requires_grad:
                data = data.detach()
            if data.is_cuda:
                data = data.cpu()
            return data.numpy()

        # Already numpy array
        if isinstance(data, np.ndarray):
            return data

        # Scalar or sequence
        if isinstance(data, (int, float, list, tuple, np.number)):
            return np.array(data)

        # Other iterables
        if isinstance(data, Iterable):
            return np.array(list(data))

        raise TypeError(f"Unsupported type for numpy conversion: {type(data)}")

    except Exception as e:
        logger.error(f"Failed to convert {type(data)} to numpy array: {e}")
        raise


def TO_SCALAR(data: TENSOR_LIKE) -> NUMERIC_SCALAR:
    """
    Convert tensor or array to Python scalar.

    Args:
        data: Input data (should contain single element)

    Returns:
        Python scalar (int or float)

    Raises:
        ValueError: If data contains more than one element

    Examples:
        >>> TO_SCALAR(torch.tensor(3.14))
        >>> TO_SCALAR(np.array([42]))
    """
    if data is None:
        raise ValueError("Cannot convert None to scalar")

    try:
        # PyTorch tensor
        if torch.is_tensor(data):
            if data.numel() != 1:
                raise ValueError(
                    f"Tensor must contain exactly one element, got {data.numel()}"
                )
            return data.item()

        # NumPy array
        if isinstance(data, np.ndarray):
            if data.size != 1:
                raise ValueError(
                    f"Array must contain exactly one element, got {data.size}"
                )
            return data.item()

        # Already scalar
        if isinstance(data, (int, float, np.number)):
            return float(data) if isinstance(data, (float, np.floating)) else int(data)

        # Single-element sequences
        if isinstance(data, (list, tuple)) and len(data) == 1:
            return TO_SCALAR(data[0])

        raise TypeError(f"Cannot convert {type(data)} to scalar")

    except Exception as e:
        logger.error(f"Failed to convert {type(data)} to scalar: {e}")
        raise


def TO_LIST(data: ARRAY_LIKE) -> List:
    """
    Convert tensor or array to Python list.

    Args:
        data: Input tensor or array

    Returns:
        Python list

    Examples:
        >>> TO_LIST(torch.tensor([1, 2, 3]))
        >>> TO_LIST(np.array([[1, 2], [3, 4]]))
    """
    if data is None:
        raise ValueError("Cannot convert None to list")

    try:
        # PyTorch tensor
        if torch.is_tensor(data):
            if data.requires_grad:
                data = data.detach()
            if data.is_cuda:
                data = data.cpu()
            return data.tolist()

        # NumPy array
        if isinstance(data, np.ndarray):
            return data.tolist()

        # Already a list
        if isinstance(data, list):
            return data

        # Tuple or other sequences
        if isinstance(data, (tuple, Sequence)):
            return list(data)

        raise TypeError(f"Cannot convert {type(data)} to list")

    except Exception as e:
        logger.error(f"Failed to convert {type(data)} to list: {e}")
        raise


def TO_DEVICE(
    data: torch.Tensor, device: Union[str, torch.device], non_blocking: bool = False
) -> torch.Tensor:
    """
    Move tensor to specified device.

    Args:
        data: Input tensor
        device: Target device
        non_blocking: Whether to use non-blocking transfer

    Returns:
        Tensor on target device

    Examples:
        >>> TO_DEVICE(tensor, 'cuda:0')
        >>> TO_DEVICE(tensor, torch.device('cpu'))
    """
    if not torch.is_tensor(data):
        raise TypeError(f"Expected tensor, got {type(data)}")

    try:
        target_device = torch.device(device)
        if data.device == target_device:
            return data
        return data.to(device=target_device, non_blocking=non_blocking)

    except Exception as e:
        logger.error(f"Failed to move tensor to device {device}: {e}")
        raise


def EXTRACT_TENSOR_FROM_INPUT(inputs: Any) -> torch.Tensor:
    """
    Extract first tensor from complex input structures.

    This function recursively searches through nested structures
    to find the first PyTorch tensor, following patterns used in
    PyTorch Lightning and HuggingFace frameworks.

    Args:
        inputs: Input of any type (tensor, tuple, list, etc.)

    Returns:
        First tensor found or None

    Examples:
        >>> EXTRACT_TENSOR_FROM_INPUT(torch.randn(2, 3))
        >>> EXTRACT_TENSOR_FROM_INPUT((torch.randn(2, 3), torch.randn(2, 4)))
        >>> EXTRACT_TENSOR_FROM_INPUT([torch.randn(2, 3), "metadata"])
    """
    if inputs is None:
        raise ValueError("Cannot extract tensor from None")

    # Direct tensor
    if torch.is_tensor(inputs):
        return inputs

    # Tuple or list - recursively search for first tensor
    if isinstance(inputs, (tuple, list)):
        for item in inputs:
            tensor = EXTRACT_TENSOR_FROM_INPUT(item)
            if tensor is not None:
                return tensor

    # Named tuple or similar structures
    elif hasattr(inputs, "_fields"):
        for field_value in inputs:
            tensor = EXTRACT_TENSOR_FROM_INPUT(field_value)
            if tensor is not None:
                return tensor

    raise ValueError("Cannot extract tensor from None")


def EXTRACT_TENSOR_FROM_OUTPUT(output: Any) -> Optional[torch.Tensor]:
    """
    Extract main tensor from model output structures.

    This function handles common output formats from deep learning models,
    including HuggingFace ModelOutput objects.

    Args:
        output: Model output of any type

    Returns:
        Main tensor output or None

    Examples:
        >>> EXTRACT_TENSOR_FROM_OUTPUT(torch.randn(2, 10))
        >>> EXTRACT_TENSOR_FROM_OUTPUT((torch.randn(2, 10), torch.randn(2, 5)))
    """
    if output is None:
        return None

    # Direct tensor
    if torch.is_tensor(output):
        return output

    # Objects with common tensor attributes (HuggingFace ModelOutput pattern)
    if hasattr(output, "last_hidden_state"):
        return output.last_hidden_state
    elif hasattr(output, "logits"):
        return output.logits
    elif hasattr(output, "hidden_states") and output.hidden_states:
        # Return last hidden state for sequence outputs
        if isinstance(output.hidden_states, (tuple, list)):
            return output.hidden_states[-1]
        return output.hidden_states

    # Tuple/List - usually first element is main output
    elif isinstance(output, (tuple, list)) and len(output) > 0:
        return EXTRACT_TENSOR_FROM_OUTPUT(output[0])

    # Objects with tensor attributes
    elif hasattr(output, "__dict__"):
        priority_attrs = ["outputs", "x", "data", "tensor"]
        for attr_name in priority_attrs:
            if hasattr(output, attr_name):
                attr_value = getattr(output, attr_name)
                tensor = EXTRACT_TENSOR_FROM_OUTPUT(attr_value)
                if tensor is not None:
                    return tensor

        # Fallback to first tensor attribute
        for attr_value in output.__dict__.values():
            if torch.is_tensor(attr_value):
                return attr_value

    return None


def VALIDATE_TENSOR_SHAPE(
    tensor: torch.Tensor, expected_shape: Tuple[int, ...], allow_batch_dim: bool = True
) -> bool:
    """
    Validate tensor shape against expected dimensions.

    Args:
        tensor: Input tensor to validate
        expected_shape: Expected shape (excluding batch dimension if allowed)
        allow_batch_dim: Whether to ignore the first dimension as batch size

    Returns:
        True if shape is valid

    Raises:
        ValueError: If shape validation fails

    Examples:
        >>> VALIDATE_TENSOR_SHAPE(torch.randn(32, 768), (768,), allow_batch_dim=True)
        >>> VALIDATE_TENSOR_SHAPE(torch.randn(224, 224, 3), (224, 224, 3))
    """
    if not torch.is_tensor(tensor):
        raise TypeError(f"Expected tensor, got {type(tensor)}")

    actual_shape = tensor.shape

    if allow_batch_dim and len(actual_shape) > len(expected_shape):
        # Compare shape excluding batch dimension
        shape_to_check = actual_shape[1:]
    else:
        shape_to_check = actual_shape

    if shape_to_check != expected_shape:
        raise ValueError(
            f"Shape mismatch: expected {expected_shape}, "
            f"got {shape_to_check} (full shape: {actual_shape})"
        )

    return True


# Utility functions for common conversion patterns
def ENSURE_TENSOR(data: TENSOR_LIKE, **kwargs) -> torch.Tensor:
    """Ensure data is a tensor, convert if necessary."""
    return TO_TENSOR(data, **kwargs) if not torch.is_tensor(data) else data


def ENSURE_NUMPY(data: TENSOR_LIKE) -> np.ndarray:
    """Ensure data is a numpy array, convert if necessary."""
    return TO_NUMPY(data) if not isinstance(data, np.ndarray) else data


def IS_TENSOR_LIKE(data: Any) -> bool:
    """Check if data can be converted to tensor."""
    try:
        TO_TENSOR(data)
        return True
    except (TypeError, ValueError):
        return False


# Export all conversion functions
__all__ = [
    "TO_TENSOR",
    "TO_NUMPY",
    "TO_SCALAR",
    "TO_LIST",
    "TO_DEVICE",
    "EXTRACT_TENSOR_FROM_INPUT",
    "EXTRACT_TENSOR_FROM_OUTPUT",
    "VALIDATE_TENSOR_SHAPE",
    "ENSURE_TENSOR",
    "ENSURE_NUMPY",
    "IS_TENSOR_LIKE",
    # Type aliases
    "TENSOR_LIKE",
    "TENSOR_INPUT",
    "MODEL_OUTPUT",
    "NUMERIC_SCALAR",
    "ARRAY_LIKE",
]
